import json
from typing import Dict, Any, Optional
from enum import Enum

class MessageType(Enum):
    """消息类型"""
    LOGIN = "LOGIN"
    LOGOUT = "LOGOUT"
    SQL_QUERY = "SQL_QUERY"
    RESPONSE = "RESPONSE"
    ERROR = "ERROR"
    HEARTBEAT = "HEARTBEAT"
    ADMIN_COMMAND = "ADMIN_COMMAND"  # 新增管理员命令类型

class Protocol:
    """通信协议"""
    
    @staticmethod
    def encode_message(msg_type: MessageType, data: Dict[str, Any], session_id: Optional[str] = None) -> bytes:
        """编码消息"""
        message = {
            "type": msg_type.value,
            "session_id": session_id,
            "data": data,
            "timestamp": None  # 可以添加时间戳
        }
        json_str = json.dumps(message, ensure_ascii=False)
        # 添加消息长度前缀（4字节）
        length = len(json_str.encode('utf-8'))
        return length.to_bytes(4, byteorder='big') + json_str.encode('utf-8')
    
    @staticmethod
    def decode_message(data: bytes) -> Optional[Dict[str, Any]]:
        """解码消息"""
        try:
            if len(data) < 4:
                return None
            
            # 读取消息长度
            length = int.from_bytes(data[:4], byteorder='big')
            if len(data) < 4 + length:
                return None
            
            # 解析JSON消息
            json_str = data[4:4+length].decode('utf-8')
            return json.loads(json_str)
        except Exception as e:
            print(f"解码消息失败: {e}")
            return None
    
    @staticmethod
    def create_login_request(username: str, password: str) -> bytes:
        """创建登录请求"""
        return Protocol.encode_message(MessageType.LOGIN, {
            "username": username,
            "password": password
        })
    
    @staticmethod
    def create_sql_request(sql: str, session_id: str) -> bytes:
        """创建SQL查询请求"""
        return Protocol.encode_message(MessageType.SQL_QUERY, {
            "sql": sql
        }, session_id)
    
    @staticmethod
    def create_response(success: bool, data: Any = None, message: str = "", session_id: Optional[str] = None) -> bytes:
        """创建响应消息"""
        return Protocol.encode_message(MessageType.RESPONSE, {
            "success": success,
            "data": data,
            "message": message
        }, session_id)
    
    @staticmethod
    def create_error(error_code: int, error_message: str, session_id: Optional[str] = None) -> bytes:
        """创建错误消息"""
        return Protocol.encode_message(MessageType.ERROR, {
            "error_code": error_code,
            "error_message": error_message
        }, session_id)